import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from typing import List, Dict, Any
#时空图嵌入服务 (TSGE)
class TemporalSpatialGraphEmbedding:
    def __init__(self, embedding_dim=128, learning_rate=0.01, epochs=100):
        self.embedding_dim = embedding_dim
        self.learning_rate = learning_rate
        self.epochs = epochs
        self.node_embeddings = {}
        self.type_embeddings = {}
        self.model = None
        self.node_to_idx = {}
        self.type_to_idx = {}
    
    class TSGEModel(nn.Module):
        def __init__(self, num_nodes, num_types, embedding_dim):
            super().__init__()
            self.node_embeddings = nn.Embedding(num_nodes, embedding_dim)
            self.type_embeddings = nn.Embedding(num_types, embedding_dim)
            self.temporal_transform = nn.Linear(embedding_dim*2, embedding_dim)
            self.edge_predictor = nn.Sequential(
                nn.Linear(embedding_dim*2, embedding_dim),
                nn.ReLU(),
                nn.Linear(embedding_dim, 1),
                nn.Sigmoid()
            )
            
        def forward(self, source_nodes, target_nodes, source_types, target_types, time_diffs):
            # 获取节点和类型嵌入
            source_emb = self.node_embeddings(source_nodes)
            target_emb = self.node_embeddings(target_nodes)
            source_type_emb = self.type_embeddings(source_types)
            target_type_emb = self.type_embeddings(target_types)
            
            # 组合节点和类型嵌入
            source_combined = source_emb + source_type_emb
            target_combined = target_emb + target_type_emb
            
            # 应用时间因素
            time_factors = torch.exp(-torch.abs(time_diffs).unsqueeze(1))
            temporal_emb = self.temporal_transform(torch.cat([source_combined, target_combined], dim=1))
            temporal_emb = temporal_emb * time_factors
            
            # 预测边概率
            edge_input = torch.cat([source_combined, target_combined], dim=1)
            return self.edge_predictor(edge_input).squeeze()
    
    def preprocess_graph_sequence(self, graph_sequence: List[Dict[str, Any]]):
        """预处理图序列数据"""
        nodes = set()
        node_types = set()
        
        # 首先收集所有节点和类型
        for graph in graph_sequence:
            for node in graph['nodes']:
                node_id = str(node['id'])
                node_type = str(node['type'])
                nodes.add(node_id)
                node_types.add(node_type)
        
        # 创建映射字典
        self.node_to_idx = {str(node): idx for idx, node in enumerate(sorted(nodes))}
        self.type_to_idx = {str(t): idx for idx, t in enumerate(sorted(node_types))}
        
        return len(nodes), len(node_types)

    def create_edge_data(self, graph_sequence: List[Dict[str, Any]]):
        """创建边数据"""
        edge_data = []
        
        for t, graph in enumerate(graph_sequence):
            # 创建节点类型查找字典
            node_type_map = {str(node['id']): str(node['type']) for node in graph['nodes']}
            
            # 处理边
            for edge in graph['edges']:
                source = str(edge['source'])
                target = str(edge['target'])
                
                # 获取节点类型
                source_type = node_type_map[source]
                target_type = node_type_map[target]
                
                # 添加正样本
                edge_data.append((
                    self.node_to_idx[source],
                    self.node_to_idx[target],
                    self.type_to_idx[source_type],
                    self.type_to_idx[target_type],
                    float(t),
                    1.0
                ))
                
                # 负采样
                nodes = list(self.node_to_idx.keys())
                for _ in range(5):
                    neg_v = np.random.choice(nodes)
                    neg_v_type = node_type_map.get(neg_v, list(self.type_to_idx.keys())[0])
                    
                    # 检查是否为已存在的边
                    if not any(e['source'] == source and e['target'] == neg_v for e in graph['edges']):
                        edge_data.append((
                            self.node_to_idx[source],
                            self.node_to_idx[neg_v],
                            self.type_to_idx[source_type],
                            self.type_to_idx[neg_v_type],
                            float(t),
                            0.0
                        ))
        
        return edge_data

    def fit(self, graph_sequence: List[Dict[str, Any]]):
        """训练模型"""
        try:
            # 预处理数据
            num_nodes, num_types = self.preprocess_graph_sequence(graph_sequence)
            
            # 创建边数据
            edge_data = self.create_edge_data(graph_sequence)
            
            if not edge_data:
                raise ValueError("No valid edge data generated")
            
            # 转换为张量
            edge_tensors = [torch.LongTensor([d[i] for d in edge_data]) for i in range(5)]
            labels = torch.FloatTensor([d[5] for d in edge_data])
            
            # 初始化模型
            self.model = self.TSGEModel(num_nodes, num_types, self.embedding_dim)
            optimizer = optim.Adam(self.model.parameters(), lr=self.learning_rate)
            criterion = nn.BCELoss()
            
            # 训练模型
            for epoch in range(self.epochs):
                optimizer.zero_grad()
                pred = self.model(
                    edge_tensors[0],
                    edge_tensors[1],
                    edge_tensors[2],
                    edge_tensors[3],
                    edge_tensors[4].float()
                )
                loss = criterion(pred, labels)
                loss.backward()
                optimizer.step()
                
                if (epoch + 1) % 10 == 0:
                    print(f"Epoch {epoch+1}/{self.epochs}, Loss: {loss.item():.4f}")
            
            # 保存嵌入
            self.save_embeddings()
            
        except Exception as e:
            print(f"Training error: {str(e)}")
            raise

    def save_embeddings(self):
        """保存训练好的嵌入"""
        with torch.no_grad():
            for node, idx in self.node_to_idx.items():
                node_tensor = torch.LongTensor([idx])
                node_emb = self.model.node_embeddings(node_tensor).numpy()[0]
                self.node_embeddings[str(node)] = node_emb

    def predict_link(self, u, v, current_time, u_type=None, v_type=None):
        """预测两个节点间的链接概率"""
        try:
            u = str(u)
            v = str(v)
            if u not in self.node_to_idx or v not in self.node_to_idx:
                return 0.0
            
            u_idx = self.node_to_idx[u]
            v_idx = self.node_to_idx[v]
            
            # 处理类型
            u_type = str(u_type) if u_type else list(self.type_to_idx.keys())[0]
            v_type = str(v_type) if v_type else list(self.type_to_idx.keys())[0]
            u_type_idx = self.type_to_idx.get(u_type, 0)
            v_type_idx = self.type_to_idx.get(v_type, 0)
            
            with torch.no_grad():
                pred = self.model(
                    torch.LongTensor([u_idx]),
                    torch.LongTensor([v_idx]),
                    torch.LongTensor([u_type_idx]),
                    torch.LongTensor([v_type_idx]),
                    torch.FloatTensor([float(current_time)])
                )
                return float(pred.item())
        except Exception as e:
            print(f"Prediction error: {str(e)}")
            return 0.0
    
    def get_embedding(self, node):
        """获取节点嵌入"""
        return self.node_embeddings.get(node) 